import numpy as np
import matplotlib.pyplot as plt
import math

class PIDController:
    def __init__(self, kp, ki, kd, max_output, dt):
        self.kp = kp
        self.ki = ki
        self.kd = kd
        self.max_output = max_output
        self.dt = dt
        self.integral = 0.0
        self.prev_error = 0.0

    def update(self, setpoint, measurement):
        error = setpoint - measurement
        self.integral += error * self.dt
        derivative = (error - self.prev_error) / self.dt
        output = self.kp * error + self.ki * self.integral + self.kd * derivative
        output = np.clip(output, -self.max_output, self.max_output)
        self.prev_error = error
        return output
class PositionController:
    """位置环：生成期望推力、姿态角和偏航角"""

    def __init__(self, dt):
        # XYZ位置PID
        self.x_pid = PIDController(1.2, 0.05, 0.8, 2.0, dt)  # 输出俯仰角
        self.y_pid = PIDController(1.2, 0.05, 0.8, 2.0, dt)  # 输出横滚角
        self.z_pid = PIDController(15.0, 2.0, 5.0, 20.0, dt)  # 输出总推力

    def update(self, pos_set, pos_meas, vel_meas):
        # 水平位置控制转换为姿态角指令（弧度）
        theta_d = self.x_pid.update(pos_set[0], pos_meas[0])
        phi_d = -self.y_pid.update(pos_set[1], pos_meas[1])

        # 高度控制（需要重力补偿）
        f_d = self.z_pid.update(pos_set[2], pos_meas[2]) + 9.81  # 重力补偿

        return phi_d, theta_d, f_d  # 返回姿态角和总推力


class AttitudeController:
    """姿态环：生成期望力矩"""

    def __init__(self, dt):
        self.phi_pid = PIDController(8.0, 0.1, 2.5, 5.0, dt)  # 横滚
        self.theta_pid = PIDController(8.0, 0.1, 2.5, 5.0, dt)  # 俯仰
        self.psi_pid = PIDController(4.0, 0.05, 1.0, 2.0, dt)  # 偏航

    def update(self, att_set, att_meas, rates_meas):
        # att_set: [phi_d, theta_d, psi_d]
        tau_phi = self.phi_pid.update(att_set[0], att_meas[0])
        tau_theta = self.theta_pid.update(att_set[1], att_meas[1])
        tau_psi = self.psi_pid.update(att_set[2], att_meas[2])
        return np.array([tau_phi, tau_theta, tau_psi])

def fM2u(f, M):
    mat = np.array([[4.179446268,       4.179446268,        4.179446268,        4.179446268],
                    [-0.6723341164784,  0.6723341164784,    0.6723341164784,    -0.6723341164784],
                    [0.6723341164784,   -0.6723341164784,   0.6723341164784,    -0.6723341164784],
                    [0.055562,          0.055562,           -0.055562,          -0.055562]])
    fM = np.vstack([f, M])
    u = np.dot(np.linalg.inv(mat), fM)
    u1 = u[0, 0]
    u2 = u[1, 0]
    u3 = u[2, 0]
    u4 = u[3, 0]
    return u1, u2, u3, u4

# 欧拉角到旋转矩阵的转换
def angle2R(roll, pitch, yaw):
    sphi = math.sin(roll)
    cphi = math.cos(roll)
    stheta = math.sin(pitch)
    ctheta = math.cos(pitch)
    spsi = math.sin(yaw)
    cpsi = math.cos(yaw)
    R = np.array([[ctheta*cpsi, sphi*stheta*cpsi-cphi*spsi, cphi*stheta*cpsi+sphi*spsi],
                  [ctheta*spsi, sphi*stheta*spsi+cphi*cpsi, cphi*stheta*spsi-sphi*cpsi],
                  [-stheta,     sphi*ctheta,                cphi*ctheta]])
    return R